
import torch
from torch.nn import Parameter
from torch_scatter import scatter_add
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax
from torch_geometric.nn.inits import glorot, zeros
import pdb

class GCNConv(MessagePassing):

    # GCNConv 的构造函数，定义了该图卷积层的各种参数：
    # in_channels：输入特征的通道数。
    # out_channels：输出特征的通道数。
    # improved：一个布尔值，指示是否应用改进的 GCN 算法。默认为 False。
    # cached：一个布尔值，指示是否缓存计算结果以加速反向传播。默认为 False。
    # bias：一个布尔值，指示是否使用偏置。默认为 True。
    # edge_norm：一个布尔值，指示是否对边进行归一化。默认为 True。
    # gfn：一个布尔值，指示是否使用 Gated Fusion Normalization (GFN)。默认为 False。
    def __init__(self,
                 in_channels,
                 out_channels,
                 improved=False,
                 cached=False,
                 bias=True,
                 edge_norm=True,
                 gfn=False):
        # 调用父类的构造函数，指定消息聚合的方式为 "add"，即将输入的消息进行求和。
        super(GCNConv, self).__init__('add')
        # 初始化各种参数，包括输入通道数、输出通道数、是否改进、是否缓存、是否使用偏置等。这些参数会以 PyTorch 的 Parameter 类型保存，并在初始化时随机初始化其值。
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.improved = improved
        self.cached = cached
        self.cached_result = None
        self.edge_norm = edge_norm
        self.gfn = gfn
        self.message_mask = None
        # 定义可学习的权重矩阵
        self.weight = Parameter(torch.Tensor(in_channels, out_channels))

        # 定义可学习的偏置向量
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()
    # reset_parameters 方法用于初始化参数。在这里，调用了 glorot 和 zeros 方法，分别对权重和偏置进行 Xavier 初始化和零初始化。
    def reset_parameters(self):
        glorot(self.weight)
        # 将偏置向量初始化为零
        zeros(self.bias)
        self.cached_result = None

    # norm 是一个静态方法，用于计算边的归一化权重。
    @staticmethod
    def norm(edge_index, num_nodes, edge_weight, improved=False, dtype=None):
        # 如果未提供边权重 edge_weight，则默认将其设置为所有边都为 1 的张量。
        if edge_weight is None:
            edge_weight = torch.ones((edge_index.size(1), ),
                                     dtype=dtype,
                                     device=edge_index.device)
        # 将边权重展平为一维张量。
        edge_weight = edge_weight.view(-1)
        
        # 断言边权重的长度与边索引的列数相同，确保数据的一致性。
        assert edge_weight.size(0) == edge_index.size(1)

        # 移除自环并添加自环。这是为了确保图的自环性质，使得每个节点都有一条指向自己的边。
        edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)
        edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
        # Add edge_weight for loop edges.
        # 为自环边添加边权重。如果 improved 为 False，则将自环的边权重设置为 1；如果 improved 为 True，则将自环的边权重设置为 2。
        loop_weight = torch.full((num_nodes, ),
                                 1 if not improved else 2,
                                 dtype=edge_weight.dtype,
                                 device=edge_weight.device)
        # 将自环的边权重添加到原有的边权重张量中
        # 是将 edge_weight 和 loop_weight 沿着指定的维度 dim=0 进行拼接，这里的 dim=0 表示在张量的第一个维度上进行拼接，即按行进行拼接。
        # 拼接后得到的张量赋值给 edge_weight，因此 edge_weight 中现在包含了原始的边权重信息以及自环边的权重信息。
        # 这样做的目的是为了在计算图卷积时考虑到自环的影响，确保每个节点都包含自身的特征信息。
        edge_weight = torch.cat([edge_weight, loop_weight], dim=0)

        # 使用 scatter_add 函数对边权重进行归一化，得到每个节点的度数。scatter_add 函数将边权重根据边的源节点进行累加，得到每个节点的度数。
        row, col = edge_index
        deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
        # 计算度数的负平方根，即将每个节点的度数取倒数再开方。这是 GCN 中常用的度数归一化操作。需要注意的是，如果某些节点的度数为 0，则其负平方根为无穷大，这可能会导致数值不稳定。因此，这里将无穷大的值替换为 0。
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        # 返回归一化后的边权重，即乘以度数的负平方根。
        return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]

    def forward(self, x, edge_index, edge_weight=None):
        """"""
        # 将输入特征 x 与权重矩阵 self.weight 进行矩阵乘法操作，以进行特征变换。
        # 矩阵乘法将输入特征 x 与权重矩阵 self.weight 相乘，得到输出特征，其中每个元素都是由输入特征和权重矩阵的对应元素相乘后的和
        x = torch.matmul(x, self.weight)
        # 如果 self.gfn 为 True，则直接返回特征变换后的结果 x，表示该图卷积层只是用于特征变换，而不进行消息传递。
        # 这个参数可以用来控制图卷积层的行为，使其符合不同的需求。
        if self.gfn:
            return x
        # 检查是否需要重新计算归一化的边权重。如果 self.cached 为 False，或者 self.cached_result 为 None，则说明需要重新计算归一化的边权重
        if not self.cached or self.cached_result is None:
            # 如果需要重新计算归一化的边权重，则根据 self.edge_norm 参数决定是否进行归一化。
            # 如果 self.edge_norm 为 True，则调用 GCNConv.norm 函数对边权重进行归一化；否则将 norm 设置为 None。
            # 然后将计算得到的边索引和归一化后的边权重缓存起来，以备后续使用。
            if self.edge_norm:
                edge_index, norm = GCNConv.norm(
                    edge_index, 
                    x.size(0), 
                    edge_weight, 
                    self.improved, 
                    x.dtype)
            else:
                norm = None
            self.cached_result = edge_index, norm
        # 从缓存中获取归一化的边权重
        edge_index, norm = self.cached_result
        # 调用 propagate 方法，执行消息传递操作。将归一化后的边权重以及输入特征 x 传入 propagate 方法中，以便在消息传递过程中使用。
        return self.propagate(edge_index, x=x, norm=norm)

    # 它用于定义消息传递过程中的消息函数。在 GCN 中，消息函数是将节点特征沿着边传递的函数。
    def message(self, x_j, norm):
        # 如果 self.edge_norm 为 True，则表示需要对消息进行归一化处理。
        # 在这种情况下，将每个节点的特征 x_j 乘以归一化后的边权重 norm，以实现消息的归一化传递。
        # 这里使用了 PyTorch 的广播机制，将 norm 张量的形状从 (num_edges,) 转换为 (num_edges, 1)，然后将其与 x_j 相乘。
        if self.edge_norm:
            return norm.view(-1, 1) * x_j
        else:
            return x_j
    # 这是 update 方法的定义，它用于定义在消息传递之后如何更新节点的特征表示。
    def update(self, aggr_out):
        # 如果存在偏置项 self.bias，则将聚合后的输出 aggr_out 加上偏置项 self.bias。这样可以为每个节点添加一个偏置项，有助于模型更好地拟合数据。
        if self.bias is not None:
            aggr_out = aggr_out + self.bias
        return aggr_out
    # 这是 __repr__ 方法的定义，它用于定义类的字符串表示形式
    def __repr__(self):
        # 此处使用 format 方法构建了一个字符串，字符串中包含了类的名称、输入特征的维度 in_channels 和输出特征的维度 out_channels。
        # 最终返回的字符串表示了这个类的基本信息，以便在打印对象时显示
        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
                                   self.out_channels)